import hashlib
import os
import json
import time
from asyncio import Semaphore
from pathlib import Path
from urllib.parse import unquote, urlparse
from aiohttp import web
import aiohttp
from typing import Optional, Dict, Any
from .utils.task_manager import task_manager
import folder_paths
import asyncio

try:
    import aiofiles
    from aiofiles import os as aio_os  # 异步os操作
except ImportError:
    raise ImportError("请安装aiofiles库: pip install aiofiles")


class CivitaiService:
    @staticmethod
    async def file_exists(file_path: Path | str) -> bool:
        """优化后的异步检查文件是否存在（兼容str和Path类型）"""
        try:
            # 先将输入转换为Path（无论原始类型是str还是Path）
            path = Path(file_path) if isinstance(file_path, str) else file_path
            return await asyncio.to_thread(path.exists)
        except Exception:
            # 发生异常时默认返回不存在（例如权限不足）
            return False

    @classmethod
    async def download_cover_if_needed(cls, images, target_cover_path, replace_cover, loop):
        """
        独立的封面下载处理方法
        负责检查封面是否需要下载并执行下载逻辑
        """
        if not images or not isinstance(images, list):
            return {
                "success": False,
                "message": "没有有效的图片数据用于封面下载"
            }

        # 筛选有效的图片格式封面URL
        valid_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.gif'}
        cover_url = ""
        for image in images:
            if 'url' in image and isinstance(image['url'], str):
                ext = Path(urlparse(image['url']).path).suffix.lower()
                if ext in valid_extensions:
                    cover_url = image['url']
                    break

        if not cover_url:
            return {
                "success": False,
                "message": "未找到有效的图片格式封面URL"
            }

        # 检查是否需要下载（不存在或强制替换）
        cover_exists = await cls.file_exists(str(target_cover_path))
        if not replace_cover and cover_exists:
            return {
                "success": True,
                "message": "封面已存在且无需替换，跳过下载"
            }

        # 执行下载
        try:
            await cls.download_and_save_cover(cover_url, str(target_cover_path), loop)
            return {
                "success": True,
                "message": "封面下载并保存成功"
            }
        except Exception as e:
            return {
                "success": False,
                "message": f"封面下载失败: {str(e)}"
            }

    @classmethod
    async def calculate_sha256(cls, file_path: str, chunk_size: int = 512 * 1024) -> str:
        lora_path = Path(file_path)
        possible_json_paths = [
            lora_path.with_suffix('.json'),  # 常规元数据文件
            lora_path.with_suffix('.marvelous-metadata.json'),  # 扩展元数据文件
        ]
        for json_path in possible_json_paths:
            if await cls.file_exists(json_path):
                try:
                    # 读取JSON文件
                    async with aiofiles.open(json_path, 'r', encoding='utf-8') as f:
                        json_content = await f.read()
                    metadata = json.loads(json_content)

                    # 从JSON中提取SHA256（兼容不同字段名）
                    sha256_value = None
                    if 'sha256' in metadata:
                        sha256_value = metadata['sha256']
                        return sha256_value
                except Exception as e:
                    print(f"[SHA256] 读取{json_path.name}失败: {str(e)}，将重新计算")
                    continue  # 读取失败则尝试下一个文件

        # 计算文件哈希
        sha256_hash = hashlib.sha256()
        async with aiofiles.open(file_path, 'rb') as f:
            while chunk := await f.read(chunk_size):
                sha256_hash.update(chunk)
        return sha256_hash.hexdigest()

    @staticmethod
    async def download_and_save_cover(url: str, save_path: str, loop) -> None:
        """下载并保存封面图片"""
        try:
            # 获取代理配置
            proxy = CivitaiService.get_comfyui_proxy()

            async with aiohttp.ClientSession() as session:
                async with session.get(url, proxy=proxy, timeout=15) as response:
                    if response.status != 200:
                        raise Exception(f"下载封面失败，状态码: {response.status}")

                    # 读取图片数据
                    image_data = await response.read()

                    # 保存到临时文件
                    temp_path = f"{save_path}.tmp"
                    async with aiofiles.open(temp_path, "wb") as f:
                        await f.write(image_data)

                    # 原子性重命名
                    await loop.run_in_executor(None, lambda: os.replace(temp_path, save_path))

        except Exception as e:
            raise Exception(f"封面处理失败: {str(e)}")

    @classmethod
    async def handle_trained_words(cls, api_data: dict, preset_path: str) -> None:
        """处理训练关键词并更新到preset文件"""
        try:
            if not await cls.file_exists(preset_path):
                return

            async with aiofiles.open(preset_path, "r", encoding="utf-8") as f:
                preset_content = await f.read()
            preset_data = json.loads(preset_content)

            # 提取训练关键词
            trained_words = api_data.get("trainedWords", [])
            if trained_words and isinstance(trained_words, list):
                # 只保留非空字符串
                valid_words = [word.strip() for word in trained_words if isinstance(word, str) and word.strip()]
                if valid_words:
                    preset_data["activation text"] = ", ".join(valid_words)

                    # 保存更新后的preset文件
                    async with aiofiles.open(preset_path, "w", encoding="utf-8") as f:
                        await f.write(json.dumps(preset_data, indent=2, ensure_ascii=False))

        except Exception as e:
            print(f"处理训练关键词失败: {str(e)}")

    @classmethod
    async def cleanup_temp_file(cls, path: str) -> None:
        """清理临时文件"""
        try:
            if await cls.file_exists(path):
                await aiofiles.os.remove(path)
        except Exception as e:
            print(f"清理临时文件失败: {str(e)}")

    @staticmethod
    def get_comfyui_proxy() -> Optional[str]:
        """
        仅针对Windows系统自动检测代理（适配主要用户群体）
        优先级：环境变量 > Windows系统代理配置
        """
        # 1. 优先读取环境变量（支持用户手动设置）
        proxy = os.getenv("HTTPS_PROXY") or os.getenv("https_proxy")
        if proxy:
            # 确保代理包含协议前缀（http:// 或 socks5://）
            if not proxy.startswith(("http://", "https://", "socks5://")):
                return f"http://{proxy}"
            return proxy

        proxy = os.getenv("HTTP_PROXY") or os.getenv("http_proxy")
        if proxy:
            if not proxy.startswith(("http://", "https://", "socks5://")):
                return f"http://{proxy}"
            return proxy

        # 2. 仅处理Windows系统代理配置
        try:
            import winreg  # Windows特有的注册表访问
            with winreg.OpenKey(
                    winreg.HKEY_CURRENT_USER,
                    r"Software\Microsoft\Windows\CurrentVersion\Internet Settings"
            ) as key:
                # 检查是否启用系统代理
                proxy_enabled, _ = winreg.QueryValueEx(key, "ProxyEnable")
                if not proxy_enabled:
                    return None  # 系统代理未启用

                # 获取代理服务器配置
                proxy_server, _ = winreg.QueryValueEx(key, "ProxyServer")
                # 解析格式："http=127.0.0.1:7890;https=127.0.0.1:7890" 或直接是 "127.0.0.1:7890"
                proxy_parts = proxy_server.split(";")
                https_proxy = None
                for part in proxy_parts:
                    if part.startswith("https="):
                        https_proxy = part.split("=")[1]
                        break
                    elif part.startswith("http="):
                        # 找不到HTTPS代理时用HTTP代理兜底
                        https_proxy = part.split("=")[1]

                # 若没有明确的http/https代理，使用默认代理
                if not https_proxy and proxy_parts:
                    https_proxy = proxy_parts[0]

                # 补充协议前缀（默认http，用户若用socks5需手动在环境变量设置）
                if https_proxy and not https_proxy.startswith(("http://", "socks5://")):
                    return f"http://{https_proxy}"
                return https_proxy

        except Exception as e:
            # Windows系统代理配置读取失败（如无权限或未配置）
            print(f"Windows系统代理检测警告: {str(e)}")
            return None

    @classmethod
    async def fetch_civitai_info(cls, request: web.Request) -> web.Response:
        temp_image_path = None
        try:
            # 解析请求数据
            data = await request.json()
            sha256 = data.get('sha256')
            lora_path_str = data.get('loraPath')
            replace_cover = data.get('replaceCover', False)

            # 路径验证
            if not lora_path_str or not isinstance(lora_path_str, str):
                return web.json_response({
                    "success": False,
                    "message": "无效的Lora文件路径"
                })

            # 解码Lora路径
            decoded_lora_path = unquote(lora_path_str)
            lora_path = Path(decoded_lora_path)

            # 构建相关文件路径
            info_json_path = lora_path.with_suffix('.marvelous-info.json')
            target_cover_path = lora_path.parent / f"{lora_path.stem}.webp"
            preset_path = lora_path.with_suffix('.json')

            # 第一步：获取信息数据（本地缓存或API）
            info_data = None
            info_source = ""
            try:
                # 尝试从本地缓存获取
                if await cls.file_exists(str(info_json_path)):
                    async with aiofiles.open(info_json_path, 'r', encoding='utf-8') as f:
                        info_data = json.loads(await f.read())
                    info_source = "本地缓存"
                else:
                    # 从API获取（处理SHA256）
                    target_sha256 = sha256 if isinstance(sha256, str) and sha256.strip() else None

                    if not target_sha256:
                        # 尝试从preset文件获取
                        try:
                            if await cls.file_exists(str(preset_path)):
                                async with aiofiles.open(preset_path, 'r', encoding='utf-8') as f:
                                    preset_data = json.loads(await f.read())
                                target_sha256 = preset_data.get('sha256') if isinstance(preset_data.get('sha256'),
                                                                                        str) else None
                        except Exception:
                            pass

                        # 计算SHA256
                        if not target_sha256:
                            if not await cls.file_exists(str(lora_path)):
                                return web.json_response({
                                    "success": False,
                                    "message": f"Lora文件不存在: {lora_path.name}"
                                })
                            target_sha256 = await cls.calculate_sha256(str(lora_path))

                            # 保存SHA256到preset文件
                            try:
                                preset_data = {}
                                if await cls.file_exists(str(preset_path)):
                                    async with aiofiles.open(preset_path, 'r', encoding='utf-8') as f:
                                        preset_data = json.loads(await f.read())
                                else:
                                    preset_data = {
                                        "description": "",
                                        "sd version": "SDXL",
                                        "activation text": "",
                                        "preferred weight": 1,
                                        "negative text": "",
                                        "notes": ""
                                    }
                                preset_data['sha256'] = target_sha256
                                async with aiofiles.open(preset_path, 'w', encoding='utf-8') as f:
                                    await f.write(json.dumps(preset_data, indent=2, ensure_ascii=False))
                            except Exception as e:
                                print(f"保存SHA256到preset文件失败: {str(e)}")

                    # 验证SHA256有效性
                    if not target_sha256 or not isinstance(target_sha256, str) or len(target_sha256) != 64:
                        return web.json_response({
                            "success": False,
                            "message": "无法获取有效的SHA256值，无法从CivitAI获取信息"
                        })

                    # 调用API获取数据
                    api_url = f"https://civitai.com/api/v1/model-versions/by-hash/{target_sha256}"
                    proxy = cls.get_comfyui_proxy()
                    print(f"使用代理访问CivitAI: {proxy or '无代理（可能无法访问）'}")

                    if proxy and not proxy.startswith(("http://", "socks5://")):
                        raise ValueError(
                            f"代理格式错误: {proxy}，正确格式应为 http://127.0.0.1:端口 或 socks5://127.0.0.1:端口")

                    timeout = aiohttp.ClientTimeout(total=30)
                    async with aiohttp.ClientSession(timeout=timeout) as session:
                        async with session.get(api_url, proxy=proxy) as response:
                            if response.status >= 400:
                                response_text = await response.text()
                                raise Exception(
                                    f"服务器返回错误状态码: {response.status}\n响应内容: {response_text[:200]}...")
                            info_data = await response.json()

                    info_source = "CivitAI API"
                    # 保存API结果到本地
                    async with aiofiles.open(info_json_path, 'w', encoding='utf-8') as f:
                        await f.write(json.dumps(info_data, indent=4, ensure_ascii=False))

            except Exception as e:
                return web.json_response({
                    "success": False,
                    "message": f"获取信息失败: {str(e)}"
                })

            # 第二步：独立处理封面（无论信息来源，强制执行）
            cover_result = await cls.download_cover_if_needed(
                images=info_data.get('images', []),
                target_cover_path=target_cover_path,
                replace_cover=replace_cover,
                loop=request.app.loop
            )

            # 处理训练关键词
            await cls.handle_trained_words(info_data, str(preset_path))

            # 整合结果返回
            return web.json_response({
                "success": True,
                "data": info_data,
                "message": f"信息已从{info_source}加载 | 封面处理: {cover_result['message']}"
            })

        except Exception as error:
            error_msg = str(error) if isinstance(error, Exception) else "获取CivitAI信息失败"
            return web.json_response({
                "success": False,
                "message": error_msg
            })

        finally:
            # 清理临时文件
            if temp_image_path:
                await cls.cleanup_temp_file(temp_image_path)
                await cls.cleanup_temp_file(f"{temp_image_path}.webp")

    @classmethod
    async def fetch_civitai_info_for_batch(
            cls,
            sha256: str,
            lora_path: str,
            replace_cover: bool,
            loop,
            app
    ) -> Dict[str, Any]:
        temp_image_path = None
        try:
            # 路径验证
            if not lora_path or not isinstance(lora_path, str):
                return {
                    "success": False,
                    "message": "无效的Lora文件路径",
                    "data": None
                }

            # 解码Lora路径
            decoded_lora_path = unquote(lora_path)
            lora_path_obj = Path(decoded_lora_path)

            # 构建相关文件路径
            info_json_path = lora_path_obj.with_suffix('.marvelous-info.json')
            target_cover_path = lora_path_obj.parent / f"{lora_path_obj.stem}.webp"
            preset_path = lora_path_obj.with_suffix('.json')

            # 第一步：获取信息数据（本地缓存或API）
            info_data = None
            info_source = ""
            try:
                if await cls.file_exists(str(info_json_path)):
                    async with aiofiles.open(info_json_path, 'r', encoding='utf-8') as f:
                        info_data = json.loads(await f.read())
                    info_source = "本地缓存"
                else:
                    # 处理SHA256
                    target_sha256 = sha256 if isinstance(sha256, str) and sha256.strip() else None

                    if not target_sha256:
                        # 尝试从preset文件获取
                        try:
                            if await cls.file_exists(str(preset_path)):
                                async with aiofiles.open(preset_path, 'r', encoding='utf-8') as f:
                                    preset_data = json.loads(await f.read())
                                target_sha256 = preset_data.get('sha256') if isinstance(preset_data.get('sha256'),
                                                                                        str) else None
                        except Exception:
                            pass

                        # 计算SHA256
                        if not target_sha256:
                            if not await cls.file_exists(str(lora_path_obj)):
                                return {
                                    "success": False,
                                    "message": f"Lora文件不存在: {lora_path_obj.name}",
                                    "data": None
                                }
                            target_sha256 = await cls.calculate_sha256(str(lora_path_obj))

                            # 保存SHA256到preset文件
                            try:
                                preset_data = {}
                                if await cls.file_exists(str(preset_path)):
                                    async with aiofiles.open(preset_path, 'r', encoding='utf-8') as f:
                                        preset_data = json.loads(await f.read())
                                else:
                                    preset_data = {
                                        "description": "",
                                        "sd version": "SDXL",
                                        "activation text": "",
                                        "preferred weight": 1,
                                        "negative text": "",
                                        "notes": ""
                                    }
                                preset_data['sha256'] = target_sha256
                                async with aiofiles.open(preset_path, 'w', encoding='utf-8') as f:
                                    await f.write(json.dumps(preset_data, indent=2, ensure_ascii=False))
                            except Exception as e:
                                print(f"保存SHA256到preset文件失败: {str(e)}")

                    # 验证SHA256有效性
                    if not target_sha256 or not isinstance(target_sha256, str) or len(target_sha256) != 64:
                        return {
                            "success": False,
                            "message": "无法获取有效的SHA256值，无法从CivitAI获取信息",
                            "data": None
                        }

                    # 从API获取数据
                    api_url = f"https://civitai.com/api/v1/model-versions/by-hash/{target_sha256}"
                    proxy = cls.get_comfyui_proxy()
                    print(f"使用代理访问CivitAI: {proxy or '无代理（可能无法访问）'}")

                    if proxy and not proxy.startswith(("http://", "socks5://")):
                        raise ValueError(
                            f"代理格式错误: {proxy}，正确格式应为 http://127.0.0.1:端口 或 socks5://127.0.0.1:端口")

                    timeout = aiohttp.ClientTimeout(total=30)
                    async with aiohttp.ClientSession(timeout=timeout) as session:
                        async with session.get(api_url, proxy=proxy) as response:
                            if response.status >= 400:
                                response_text = await response.text()
                                raise Exception(
                                    f"服务器返回错误状态码: {response.status}\n响应内容: {response_text[:200]}...")
                            info_data = await response.json()

                    info_source = "CivitAI API"
                    # 保存API结果到本地
                    async with aiofiles.open(info_json_path, 'w', encoding='utf-8') as f:
                        await f.write(json.dumps(info_data, indent=4, ensure_ascii=False))

            except Exception as e:
                return {"success": False, "message": f"获取信息失败: {str(e)}", "data": None}

            # 第二步：独立处理封面（无论信息来源，强制执行）
            cover_result = await cls.download_cover_if_needed(
                images=info_data.get('images', []),
                target_cover_path=target_cover_path,
                replace_cover=replace_cover,
                loop=loop
            )

            # 处理训练关键词
            await cls.handle_trained_words(info_data, str(preset_path))

            return {
                "success": True,
                "data": info_data,
                "message": f"信息已从{info_source}加载 | 封面处理: {cover_result['message']}"
            }

        except Exception as error:
            error_msg = str(error) if isinstance(error, Exception) else "获取CivitAI信息失败"
            return {
                "success": False,
                "message": error_msg,
                "data": None
            }

        finally:
            # 清理临时文件
            if temp_image_path:
                await cls.cleanup_temp_file(temp_image_path)
                await cls.cleanup_temp_file(f"{temp_image_path}.webp")

    @classmethod
    async def start_batch_marvelous_check(cls, request: web.Request) -> web.Response:
        """启动批量检查marvelous-info.json文件的任务"""
        TASK_ID = "marvelous_info_global_check"

        # 检查任务是否已存在
        if TASK_ID in task_manager.tasks:
            current_progress = task_manager.get_progress(TASK_ID)
            return web.json_response({
                "success": True,
                "message": "批量检查任务已在运行中",
                "task_id": TASK_ID,
                "progress": current_progress,
                "status": "running"
            })

        try:
            # 获取所有LoRA文件夹
            lora_folders = folder_paths.folder_names_and_paths.get("loras", ([], []))[0]
            if not lora_folders:
                return web.json_response({
                    "success": False,
                    "message": "未配置LoRA文件夹路径"
                })

            # 收集所有LoRA文件
            lora_files = []
            for folder in lora_folders:
                if not await cls.file_exists(Path(folder)):
                    continue

                async def collect_lora_files(current_path: str):
                    entries = await request.loop.run_in_executor(
                        None,
                        lambda: [(entry.name, entry.is_dir()) for entry in os.scandir(current_path)]
                    )

                    for entry_name, is_dir in entries:
                        full_path = os.path.join(current_path, entry_name)
                        if is_dir:
                            await collect_lora_files(full_path)
                        elif entry_name.endswith('.safetensors'):
                            lora_files.append(full_path)

                await collect_lora_files(folder)

            if not lora_files:
                return web.json_response({
                    "success": True,
                    "message": "未找到任何LoRA文件",
                    "task_id": TASK_ID,
                    "status": "completed"
                })

            # 初始化任务
            total_count = len(lora_files)
            task_manager.create_task(TASK_ID, total_count)

            # 注册取消函数
            def cancel_check():
                print(f"批量检查任务 {TASK_ID} 已被取消")

            task_manager.add_active_export(TASK_ID, cancel_check)

            # 后台执行批量检查
            async def background_batch_check():
                processed_count = 0
                batch_start_time = time.time()
                print(f"[{time.strftime('%H:%M:%S')}] 开始处理文件，共{total_count}个文件")

                # 控制并发数量
                max_concurrent = 5
                semaphore = Semaphore(max_concurrent)

                async def process_file(index, lora_path_str):
                    nonlocal processed_count
                    async with semaphore:
                        file_start_time = time.time()
                        lora_path = Path(lora_path_str)
                        result = {
                            "file_path": str(lora_path),
                            "file_name": lora_path.name,
                            "success": False,
                            "message": ""
                        }
                        print(
                            f"\n[{time.strftime('%H:%M:%S')}] 开始处理第{index}/{total_count}个文件: {lora_path.name}")

                        try:
                            # 即使info.json存在，也需要检查封面（通过fetch_civitai_info_for_batch实现）
                            sha256 = await cls.calculate_sha256(lora_path_str)
                            info_result = await cls.fetch_civitai_info_for_batch(
                                sha256=sha256,
                                lora_path=str(lora_path),
                                replace_cover=True,
                                loop=request.app.loop,
                                app=request.app
                            )
                            result["success"] = info_result["success"]
                            result["message"] = info_result["message"]

                        except Exception as e:
                            print(f"[{time.strftime('%H:%M:%S')}] 处理失败: {str(e)}")
                            result["message"] = f"处理失败: {str(e)}"

                        finally:
                            processed_count += 1
                            task_manager.update_progress(TASK_ID, processed_count)
                            print(
                                f"[{time.strftime('%H:%M:%S')}] 第{index}/{total_count}个文件处理完成，耗时: {time.time() - file_start_time:.2f}秒")

                        return result

                # 并发执行所有文件处理任务
                tasks = []
                for index, lora_path_str in enumerate(lora_files, 1):
                    if TASK_ID not in task_manager.tasks:
                        print(f"[{time.strftime('%H:%M:%S')}] 任务已取消，停止处理")
                        break
                    tasks.append(process_file(index, lora_path_str))

                results = await asyncio.gather(*tasks)
                await task_manager.emit('task_complete', TASK_ID, results)
                print(
                    f"\n[{time.strftime('%H:%M:%S')}] 所有文件处理完成，总耗时: {time.time() - batch_start_time:.2f}秒")

            # 启动后台任务
            asyncio.create_task(background_batch_check())

            return web.json_response({
                "success": True,
                "message": "批量检查任务已启动",
                "task_id": TASK_ID,
                "total_files": total_count,
                "status": "running"
            })

        except Exception as e:
            return web.json_response({
                "success": False,
                "message": f"启动批量检查任务失败: {str(e)}"
            })
